% 2D Time Reversal Reconstruction For A Circular Sensor  vs a Linear Sensor
%
%
% author: YANG Chunshan
% date: 7th August 2025

clearvars;


% =========================================================================
% Linear Detector SIMULATION with TR
% =========================================================================

% create the computational grid
PML_size = 20;              % size of the PML in grid points
Nx = 256 - 2 * PML_size;    % number of grid points in the x direction
Ny = 256 - 2 * PML_size;    % number of grid points in the y direction
x = 10e-3;          % total grid size [m]
y = 10e-3;          % total grid size [m]
dx = x/Nx;                % grid point spacing in the x direction [m]
dy = y/Ny;                % grid point spacing in the y direction [m]
kgrid = kWaveGrid(Nx, dx, Ny, dy);

% define the properties of the propagation medium
medium.sound_speed = 1500;	% [m/s]

% create initial pressure distribution using makeDisc
disc_magnitude = 5;         % [Pa]
disc_x_pos = 110;            % [grid points]
disc_y_pos = 140;           % [grid points]
disc_radius = 5;            % [grid points]
disc_2 = disc_magnitude * makeDisc(Nx, Ny, disc_x_pos, disc_y_pos, disc_radius);

disc_x_pos = 110;            % [grid points]
disc_y_pos = 110;           % [grid points]
disc_radius = 8;            % [grid points]
disc_1 = disc_magnitude * makeDisc(Nx, Ny, disc_x_pos, disc_y_pos, disc_radius);

% smooth the initial pressure distribution and restore the magnitude
p0 = smooth(disc_1 + disc_2, true);

% assign to the source structure
source.p0 = p0;

% define a binary line sensor
sensor.mask = zeros(Nx, Ny);
sensor.mask(1, :) = 1;

% create the time array
kgrid.makeTime(medium.sound_speed);

% set the input arguements: force the PML to be outside the computational
% grid; switch off p0 smoothing within kspaceFirstOrder2D
input_args = {'PMLInside', false, 'PMLSize', PML_size, 'Smooth', false, 'PlotPML', false};

% run the simulation
sensor_data = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

% reset the initial pressure
source.p0 = 0;

% assign the time reversal data
sensor.time_reversal_boundary_data = sensor_data;

% run the time reversal reconstruction
p0_recon = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

% add first order compensation for only recording over a half plane
p0_recon = 2 * p0_recon;

% % repeat the FFT reconstruction for comparison
% p_xy = kspaceLineRecon(sensor_data.', dy, kgrid.dt, medium.sound_speed, ...
%     'PosCond', true, 'Interp', '*linear');
% 
% % define a second k-space grid using the dimensions of p_xy
% [Nx_recon, Ny_recon] = size(p_xy);
% kgrid_recon = kWaveGrid(Nx_recon, kgrid.dt * medium.sound_speed, Ny_recon, dy);
% 
% % resample p_xy to be the same size as source.p0
% p_xy_rs = interp2(kgrid_recon.y, kgrid_recon.x - min(kgrid_recon.x(:)), p_xy, kgrid.y, kgrid.x - min(kgrid.x(:)));


% =========================================================================
% Circular Detector SIMULATION with TR
% =========================================================================

kgrid2 = kWaveGrid(Nx, dx, Ny, dy);
% define the properties of the propagation medium
medium2.sound_speed = 1500;	% [m/s]

% create initial pressure distribution using makeDisc
disc_magnitude = 5;         % [Pa]
disc_x_pos = 110;            % [grid points]
disc_y_pos = 140;           % [grid points]
disc_radius = 5;            % [grid points]
disc_2 = disc_magnitude * makeDisc(Nx, Ny, disc_x_pos, disc_y_pos, disc_radius);

disc_x_pos = 110;            % [grid points]
disc_y_pos = 110;           % [grid points]
disc_radius = 8;            % [grid points]
disc_1 = disc_magnitude * makeDisc(Nx, Ny, disc_x_pos, disc_y_pos, disc_radius);

% smooth the initial pressure distribution and restore the magnitude
p0 = smooth(disc_1 + disc_2, true);

% assign to the source structure
source2.p0 = p0;

% % define the properties of the propagation medium
% medium.sound_speed = 1500;  % [m/s]

% define a centered Cartesian circular sensor
sensor_radius = 4.5e-3;     % [m]
sensor_angle = 4*pi/2;      % [rad]
sensor_pos = [0, 0];        % [m]
num_sensor_points = 256;
cart_sensor_mask = makeCartCircle(sensor_radius, num_sensor_points, sensor_pos, sensor_angle);

% assign to sensor structure
sensor2.mask = cart_sensor_mask;

% create the time array
kgrid2.makeTime(medium2.sound_speed);

% set the input options
input_args2 = {'Smooth', false, 'PMLInside', false, 'PlotPML', false};

% run the simulation
sensor_data2 = kspaceFirstOrder2D(kgrid2, medium2, source2, sensor2, input_args2{:});

% add noise to the recorded sensor data
signal_to_noise_ratio = 40;	% [dB]
sensor_data2 = addNoise(sensor_data2, signal_to_noise_ratio, 'peak');

% create a second computation grid for the reconstruction to avoid the
% inverse crime
Nx = 300;           % number of grid points in the x direction
Ny = 300;           % number of grid points in the y direction
dx = x/Nx;          % grid point spacing in the x direction [m]
dy = y/Ny;          % grid point spacing in the y direction [m]
kgrid_recon2 = kWaveGrid(Nx, dx, Ny, dy);

% use the same time array for the reconstruction
kgrid_recon2.setTime(kgrid2.Nt, kgrid2.dt);

% reset the initial pressure
source2.p0 = 0;

% assign the time reversal data
sensor2.time_reversal_boundary_data = sensor_data2;

% run the time-reversal reconstruction
p0_recon2 = kspaceFirstOrder2D(kgrid_recon2, medium2, source2, sensor2, input_args2{:});

% % resample p0_recon2 to be the same size as source.p0
p0_recon2 = interp2(kgrid_recon2.y, kgrid_recon2.x - min(kgrid_recon2.x(:)), p0_recon2, kgrid2.y, kgrid2.x - min(kgrid2.x(:)), 'linear', 0);

% % create a binary sensor mask of an equivalent continuous circle 
% sensor_radius_grid_points = round(sensor_radius / kgrid_recon2.dx);
% binary_sensor_mask = makeCircle(kgrid_recon2.Nx, kgrid_recon2.Ny, kgrid_recon2.Nx/2 + 1, kgrid_recon2.Ny/2 + 1, sensor_radius_grid_points, sensor_angle);
% 
% % assign to sensor structure
% sensor2.mask = binary_sensor_mask;
% 
% % interpolate data to remove the gaps and assign to sensor structure
% sensor2.time_reversal_boundary_data = interpCartData(kgrid_recon2, sensor_data2, cart_sensor_mask, binary_sensor_mask);
% 
% % run the time-reversal reconstruction
% p0_recon_interp = kspaceFirstOrder2D(kgrid_recon2, medium2, source2, sensor2, input_args2{:});

% =========================================================================
% VISUALISATION
% =========================================================================

% plot the initial pressure and sensor distribution
figure;
imagesc(kgrid.y_vec * 1e3, kgrid.x_vec * 1e3, p0 + sensor.mask * disc_magnitude, [-disc_magnitude, disc_magnitude]);
colormap(getColorMap);
ylabel('x-position [mm]');
xlabel('y-position [mm]');
axis image;
colorbar;
%scaleFig(1, 0.65);

% % plot the reconstructed initial pressure 
% figure;
% imagesc(kgrid.y_vec * 1e3, kgrid.x_vec * 1e3, p0_recon, [-disc_magnitude, disc_magnitude]);
% colormap(getColorMap);
% ylabel('x-position [mm]');
% xlabel('y-position [mm]');
% axis image;
% colorbar;
% %scaleFig(1, 0.65);

% apply a positivity condition
p0_recon(p0_recon < 0) = 0;

% plot the reconstructed initial pressure with positivity condition
figure;
imagesc(kgrid.y_vec * 1e3, kgrid.x_vec * 1e3, p0_recon, [-disc_magnitude, disc_magnitude]);
colormap(getColorMap);
ylabel('x-position [mm]');
xlabel('y-position [mm]');
axis image;
colorbar;
%scaleFig(1, 0.65);

% % plot a profile for comparison
% figure;
% plot(kgrid.y_vec * 1e3, p0(disc_x_pos, :), 'k-', ...
%      kgrid.y_vec * 1e3, p0_recon(disc_x_pos, :), 'b:');
% xlabel('y-position [mm]');
% ylabel('Pressure');
% legend('Initial Pressure', 'Time Reversal');
% axis tight;
% set(gca, 'YLim', [0, 5.1]);

% plot the initial pressure and sensor distribution
figure;
imagesc(kgrid2.y_vec * 1e3, kgrid2.x_vec * 1e3, p0 + cart2grid(kgrid2, cart_sensor_mask), [-1, 1]);
colormap(getColorMap);
ylabel('x-position [mm]');
xlabel('y-position [mm]');
axis image;

% % plot the simulated sensor data
% figure;
% imagesc(sensor_data2, [-1, 1]);
% colormap(getColorMap);
% ylabel('Sensor Position');
% xlabel('Time Step');
% colorbar;

p0_recon2(p0_recon2 < 0) = 0;
% plot the reconstructed initial pressure 
figure;
imagesc(kgrid_recon2.y_vec * 1e3, kgrid_recon2.x_vec * 1e3, p0_recon2, [-disc_magnitude, disc_magnitude]);
colormap(getColorMap);
ylabel('x-position [mm]');
xlabel('y-position [mm]');
axis image;
colorbar;

% % plot the reconstructed initial pressure using the interpolated data
% figure;
% imagesc(kgrid_recon2.y_vec * 1e3, kgrid_recon2.x_vec * 1e3, p0_recon_interp, [-1, 1]);
% colormap(getColorMap);
% ylabel('x-position [mm]');
% xlabel('y-position [mm]');
% axis image;

% plot a profile for comparison
figure;
plot(kgrid2.y_vec * 1e3, p0(disc_x_pos, :), 'k--', ...
     kgrid2.y_vec * 1e3, p0_recon2(disc_x_pos, :), 'r-', ...
     kgrid.y_vec * 1e3, p0_recon(disc_x_pos, :), 'b:','linewidth', 4);
xlabel('y-position [mm]');
ylabel('Pressure');
legend('Initial Pressure', 'Circle Detector', 'Linear Detector');
axis tight;
set(gca, 'YLim', [0 5.1]);
% set(gca, 'box', 'off',...
%          'linewidth',4.5,...
%          'TickDir','in',...
%          'fontsize',16,'fontname','Times','FontWeight','bold','dataaspectratio',[1,1,1],'box','on')

% 计算PSNR和SSIM值
% 对于线性探测器重建
psnr_linear = psnr(p0_recon, p0);
ssim_linear = ssim(p0_recon, p0);

psnr_circular = psnr(p0_recon2, p0);
ssim_circular = ssim(p0_recon2, p0);

% 显示结果
fprintf('线性探测器重建:\n');
fprintf('PSNR: %.2f dB\n', psnr_linear);
fprintf('SSIM: %.4f\n\n', ssim_linear);

fprintf('圆形探测器重建:\n');
fprintf('PSNR: %.2f dB\n', psnr_circular);
fprintf('SSIM: %.4f\n', ssim_circular);